package burp;

import org.apache.commons.lang.StringEscapeUtils;

import java.io.PrintWriter;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Author: b0lu email: tasmart@163.com QQ: 10461540 布鲁's blog http://b0lu.xyz
 */
public class BurpExtender implements IBurpExtender, IHttpListener,
        IProxyListener, IScannerListener, IExtensionStateListener {
    private IBurpExtenderCallbacks _callbacks;
    private IExtensionHelpers _helpers;
    private PrintWriter _stdout;

    @Override
    public void registerExtenderCallbacks(final IBurpExtenderCallbacks callbacks) {
        // keep a reference to our callbacks object
        this._callbacks = callbacks;
        this._helpers = callbacks.getHelpers();

        // set our extension name
        callbacks.setExtensionName("b0lu unicode decode");

        // obtain our output stream
        this._stdout = new PrintWriter(callbacks.getStdout(), true);

        // register ourselves as an HTTP listener
        callbacks.registerHttpListener(this);

        // register ourselves as a Proxy listener
        callbacks.registerProxyListener(this);

        // register ourselves as a Scanner listener
        callbacks.registerScannerListener(this);

        // register ourselves as an extension state listener
        callbacks.registerExtensionStateListener(this);
    }

    @Override
    public void extensionUnloaded() {
        // TODO Auto-generated method stub
        this._stdout.println("Extension was unloaded");
    }

    @Override
    public void newScanIssue(IScanIssue issue) {
        // TODO Auto-generated method stub
        return;
    }

    @Override
    public void processProxyMessage(boolean messageIsRequest,
                                    IInterceptedProxyMessage message) {
        // TODO Auto-generated method stub
    }

    @Override
    public void processHttpMessage(int toolFlag, boolean messageIsRequest,
                                   IHttpRequestResponse messageInfo) {
        // TODO Auto-generated method stub
        IRequestInfo analyzedReq = this._helpers.analyzeRequest(messageInfo);
        if (!messageIsRequest) {

            byte[] bresponse = messageInfo.getResponse();
            IResponseInfo resp = this._helpers.analyzeResponse(bresponse);
            List<String> headers = resp.getHeaders();
            int b_body_length = bresponse.length - resp.getBodyOffset();
            byte[] bresponse_body = new byte[b_body_length];
            System.arraycopy(bresponse, resp.getBodyOffset(), bresponse_body, 0, b_body_length);

            String response_body = new String(bresponse_body);
            String convert_response_body = this.unicodeToString(response_body);

            messageInfo.setResponse(this._helpers.buildHttpMessage(headers, convert_response_body.getBytes()));
        }

    }

    /**
     * 转化unicode为汉字
     *
     * @param str
     * @return
     */
    public String unicodeToString(String str) {

        Pattern pattern = Pattern.compile("(\\\\u(\\p{XDigit}{4}))");
        Matcher matcher = pattern.matcher(str);
        char ch;
        while (matcher.find()) {
            ch = (char) Integer.parseInt(matcher.group(2), 16);
            str = str.replace(matcher.group(1), ch + "");
        }
        return StringEscapeUtils.unescapeHtml(str);
    }

}
